#!/usr/bin/env python3
import argparse

import common_modules
from intra_layer_speedup import *
from converge_fig import *
from stream_fig import *
from mem_overhead_fig import *
from time_overhead import *
from mot_fig import *
from pipeline_speedup import *

def str2bool(v):
    if isinstance(v, bool):
        return v

    if v.lower() in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected.')

if __name__ == "__main__":
    # Parse command line arguments.
    parser = argparse.ArgumentParser(description='Analyze and draw HGP4NN experiment figures.')
    parser.add_argument('-root', nargs=1, type=str, help='Root log folder.', dest='root')
    parser.add_argument('-serial', nargs=1, type=str, help='Serial log file folder.', dest='serial')
    parser.add_argument('-hgp4nn', nargs=1, type=str, help='Log file folder for the HGP4NN framework.', dest='hgp4nn')
    parser.add_argument('-output', nargs=1, type=str, help='Output directory.', dest='output')
    parser.add_argument('-intra', nargs='?', default=False, type=str2bool, help='Analyze and draw IntraLayer-related figures.', dest='intra')
    parser.add_argument('-inter', nargs='?', default=False, type=str2bool, help='Analyze and draw IntraLayer-related figures.', dest='inter')
    parser.add_argument('-mem', nargs='?', default=False, type=str2bool, help='Analyze and draw memory overhead figures.', dest='mem')
    parser.add_argument('-stream', nargs='?', default=False, type=str2bool, help='Analyze and draw the stream figure.', dest='stream')
    parser.add_argument('-time', nargs='?', default=False, type=str2bool, help='Analyze and output the time overhead results.', dest='time')
    parser.add_argument('-converge', nargs='?', default=False, type=str2bool, help='Analyze and output the converge figure.', dest='converge')
    parser.add_argument('-all', nargs='?', default=False, type=str2bool, help='Output all results.', dest='all')
    parser.add_argument('-mot', nargs='?', default=False, type=str2bool, help='Analyze and draw motivation figures.', dest='mot')
    parser.add_argument('-pipe', nargs='?', default=False, type=str2bool, help='Analyze and draw the overall speedup figures.', dest='pipe')
    parser.add_argument('-mfolder', nargs=1, type=str, help='Folder used to store the motivation data.', dest='mfolder')

    args = parser.parse_args()

    # Parse serial log files.
    # root_folder = '~/workspace/GLP4NN/Expr-HGP4NN'
    devices = ["P100", "TitanXP", "K40c"]

    if args.intra:
        print("Analyze and draw IntraLayer-related figures...")
        # net_layer_used = DrawIntraLayerSpeedup(args, devices)
        # print(net_layer_used)
        if args.mem:
            print("Analyze and draw the memory overhead figure...")
            DrawBuffOverheadFig(args, devices)
        if args.stream:
            print("Analyze and draw the number of streams...")
            DrawStreamFigs(args, devices, net_layer_used)
        if args.time:
            print("Analyze and output the time overhead figure...")
            TimeOverhead(args, devices)
        if args.converge:
            print("Analyze and draw the converge figure...")
            DrawConvergeLossLine(args)
        if args.pipe:
            print('Analyze and draw the overall speedup figures...')
            DrawOverallSpeedupFig(args, devices)
        if args.all:
            print("Analyze and draw all figures...")
            DrawBuffOverheadFig(args, devices)
            DrawStreamFigs(args, devices, net_layer_used)
            TimeOverhead(args, devices)
            DrawConvergeLossLine(args)
            DrawOverallSpeedupFig(args, devices)

    if args.inter:
        print("Aanlyze and draw InterLayer-related figures...")
        # DrawInterLayerSpeedup(args, devices)

    if args.mot:
        mot_devices = ['K40c', 'TitanXP']
        DrawMotFigs(args, mot_devices)
